from typing import Optional, Tuple, Union
import torch
from torch import Tensor

from torch.nn import Parameter
from torch_geometric.nn.conv import MessagePassing
from torch_geometric.nn.dense.linear import Linear
from torch_geometric.nn.inits import zeros
from torch_geometric.typing import (
    Adj,
    OptPairTensor,
    OptTensor,
    SparseTensor,
)
from torch_geometric.utils import add_remaining_self_loops
from torch_geometric.utils import (
    scatter,
    spmm,
)
from torch_geometric.utils.num_nodes import maybe_num_nodes


def add_remaining_self_loops(edge_index, edge_attr=None, fill_value=None, num_nodes=None):
    print("==================2、增加自循环边=======================")
    print("原始 edge_index")
    print(edge_index)
    print("原始 edge_attr")
    print(edge_attr)
    N = num_nodes
    # 为源和目的不一样的边设定mask
    mask = edge_index[0] != edge_index[1]
    loop_index = torch.arange(0, N, dtype=torch.long, device=edge_index.device)
    loop_index = loop_index.unsqueeze(0).repeat(2, 1)
    # 建立一个新的二维矩阵，用于处理源和目的一致的情况
    if edge_attr is not None:
        if isinstance(fill_value, (int, float)):
            loop_attr = edge_attr.new_full((N,) + edge_attr.size()[1:],
                                           fill_value)
        inv_mask = ~mask
        loop_attr[edge_index[0][inv_mask]] = edge_attr[inv_mask]
        edge_attr = torch.cat([edge_attr[mask], loop_attr], dim=0)

    # 在没有边属性的情况
    edge_index = torch.cat([edge_index[:, mask], loop_index], dim=1)
    print("增加自环之后的边", edge_index)
    print("增加自环之后的边权重", edge_attr)
    print("=========================================")
    return edge_index, edge_attr


def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
             add_self_loops=True, flow="source_to_target", dtype=None):
    fill_value = 2. if improved else 1.

    if add_self_loops:
        edge_index, edge_weight = add_remaining_self_loops(
            edge_index, edge_weight, fill_value, num_nodes)

    if edge_weight is None:
        # 边缘长度个1
        edge_weight = torch.ones((edge_index.size(1),), dtype=dtype,
                                 device=edge_index.device)
    print("==================3、度矩阵变换=======================")
    row, col = edge_index[0], edge_index[1]
    idx = col if flow == 'source_to_target' else row
    # 如果流向为source_to_target则对被指向的节点信息进行聚合
    print(edge_weight, idx, num_nodes)
    deg = scatter(edge_weight, idx, dim=0, dim_size=num_nodes, reduce='sum')
    print("度矩阵计算结果", deg)
    deg_inv_sqrt = deg.pow_(-0.5)
    deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0)
    print("度矩阵变换之后的内容", deg_inv_sqrt)
    # 左乘进行行变换，右乘进行列变换
    # print(row,deg_inv_sqrt[row])
    # print(col,deg_inv_sqrt[col])
    edge_weight = deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
    print("变换之后的边权重", edge_weight)
    print("=========================================")
    return edge_index, edge_weight


class GCNConv(MessagePassing):
    _cached_edge_index: Optional[OptPairTensor]
    _cached_adj_t: Optional[SparseTensor]

    def __init__(
            self,
            in_channels,
            out_channels,
            improved=False,
            cached=False,
            add_self_loops=True,
            normalize=True,
            bias=True,
            **kwargs,
    ):
        kwargs.setdefault('aggr', 'add')
        super().__init__(**kwargs)

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.improved = improved
        self.cached = cached
        self.add_self_loops = add_self_loops
        self.normalize = normalize

        self._cached_edge_index = None
        self._cached_adj_t = None

        self.lin = Linear(in_channels, out_channels, bias=False,
                          weight_initializer='glorot')

        if bias:
            self.bias = Parameter(torch.empty(out_channels))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self):
        super().reset_parameters()
        self.lin.reset_parameters()
        zeros(self.bias)
        self._cached_edge_index = None
        self._cached_adj_t = None

    def forward(self, x: Tensor, edge_index: Adj,
                edge_weight: OptTensor = None) -> Tensor:

        if self.normalize:
            cache = self._cached_edge_index
            if cache is None:
                print("==================1、第一次聚合统计信息=======================")
                print("cache is None")
                print("边信息 edge_index", edge_index)
                print("边权重 edge_weight", edge_weight)
                print("x.size(self.node_dim) 节点数", x.size(self.node_dim))
                print("是否增益权重 self.improved,", self.improved)
                print("是否增加自环 self.add_self_loops, ", self.add_self_loops)
                print("方向 self.flow", self.flow)
                print("节点特征数据类型 x.dtype", x.dtype)
                # 做 GCN Normlize
                edge_index, edge_weight = gcn_norm(  # yapf: disable
                    edge_index, edge_weight, x.size(self.node_dim),
                    self.improved, self.add_self_loops, self.flow, x.dtype)
                if self.cached:
                    self._cached_edge_index = (edge_index, edge_weight)
            else:
                edge_index, edge_weight = cache[0], cache[1]
        print("重建后的 edge_weight", edge_weight)
        print("=====================================")
        x = self.lin(x)
        # propagate_type: (x: Tensor, edge_weight: OptTensor)
        out = self.propagate(edge_index, x=x, edge_weight=edge_weight,
                             size=None)
        if self.bias is not None:
            out = out + self.bias
        return out, edge_weight

    # def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
    #     #对 x_j 进行更新 （发起端）
    #     return x_j if edge_weight is None else edge_weight.view(-1, 1) * x_j
    #
    # def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor:
    #     print("####################################################################################################################################################################################")
    #     print("######################### message_and_aggregate 处理前1")
    #     print(adj_t)
    #     print(x)
    #     print("######################### message_and_aggregate 处理前1")
    #     return spmm(adj_t, x, reduce=self.aggr)

    # def update(self,x_i):
    #     return x_i


# edge_index = torch.tensor([[0, 1, 2, 3, 4], [1, 2, 0, 4, 3]], dtype=torch.long)  # 2 x E
# x = torch.tensor([[0.1, 0.1], [0.2, 0.2], [0.3, 0.3], [0.4, 0.4], [0.5, 0.5]], dtype=torch.float)  # N x emb(in)

edge_index = torch.tensor([[0, 1, 2, 3, 4], [1, 2, 0, 4, 3]], dtype=torch.long)  # 2 x E
x = torch.tensor([[0.1, 0.1], [0.2, 0.2], [0.3, 0.3], [0.4, 0.4], [0.5, 0.5]], dtype=torch.float)
edge_weight = torch.FloatTensor([[1], [1], [1], [1], [1]])

conv1 = GCNConv(2, 2, cached=False)
# conv1 = GCNConv(2, 2,cached=True)

edge_weight_for_graph = [i[0] for i in edge_weight]
out1, edge_weight1 = conv1(x, edge_index, edge_weight)
out2, edge_weight2 = conv1(out1, edge_index, edge_weight1[0:5])
out3, edge_weight3 = conv1(out2, edge_index, edge_weight2[0:5])
out4, edge_weight4 = conv1(out3, edge_index, edge_weight3[0:5])
out5, edge_weight5 = conv1(out4, edge_index, edge_weight4[0:5])
out6, edge_weight6 = conv1(out3, edge_index, edge_weight3[0:5])

out1 = out1.detach()
out2 = out2.detach()
out3 = out3.detach()
out4 = out4.detach()
out5 = out5.detach()
out6 = out6.detach()
edge_weight_for_graph1 = [i[0] for i in edge_weight1.detach()]
edge_weight_for_graph2 = [i[0] for i in edge_weight2.detach()]
edge_weight_for_graph3 = [i[0] for i in edge_weight3.detach()]
edge_weight_for_graph4 = [i[0] for i in edge_weight4.detach()]
edge_weight_for_graph5 = [i[0] for i in edge_weight5.detach()]
edge_weight_for_graph6 = [i[0] for i in edge_weight6.detach()]

import torch
import matplotlib.pyplot as plt

import numpy as np

# plt.style.use(['ggplot', 'grayscale'])
fig = plt.figure()
lines = 4


def new_x_y(x, y, edge_weight):
    new_x = [x[i] * edge_weight[i] for i in range(len(x))]
    new_y = [y[i] * edge_weight[i] for i in range(len(y))]
    return new_x, new_y


x_1, y_1 = x.t()[0], x.t()[1]
x_2, y_2 = out1.t()[0], out1.t()[1]
x_3, y_3 = out2.t()[0], out2.t()[1]
x_4, y_4 = out3.t()[0], out3.t()[1]
x_5, y_5 = out4.t()[0], out4.t()[1]
x_6, y_6 = out5.t()[0], out5.t()[1]
x_7, y_7 = out6.t()[0], out6.t()[1]

x_1_1, y_1_1 = new_x_y(x.t()[0], x.t()[1], edge_weight)
x_2_1, y_2_1 = new_x_y(out1.t()[0], out1.t()[1], edge_weight1)
x_3_1, y_3_1 = new_x_y(out2.t()[0], out2.t()[1], edge_weight2)
x_4_1, y_4_1 = new_x_y(out3.t()[0], out3.t()[1], edge_weight3)
x_5_1, y_5_1 = new_x_y(out4.t()[0], out4.t()[1], edge_weight4)
x_6_1, y_6_1 = new_x_y(out5.t()[0], out5.t()[1], edge_weight5)
x_7_1, y_7_1 = new_x_y(out6.t()[0], out6.t()[1], edge_weight6)

colors = np.random.rand(len(x), 3)  # 生成50个RGB颜色值，每个颜色值是一个长度为3的数组
print("x1")
ax3 = fig.add_subplot(lines, 3, 1)
ax3.scatter(x_1, y_1, c=colors)
# 为每个点添加序号标签
for i, txt in enumerate(x_1):
    ax3.text(x_1[i], y_1[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')

ax3 = fig.add_subplot(lines, 3, 2)
ax3.scatter(x_1_1, y_1_1, c=colors)
# 为每个点添加序号标签
for i, txt in enumerate(x_1):
    ax3.text(x_1_1[i], y_1_1[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')

ax3 = fig.add_subplot(lines, 3, 3)
ax3.scatter([1.0, 2.0, 3.0, 4.0, 5.0], edge_weight_for_graph)
ax3.set_ylim((-0, 1))
ax3.set_xlim((0, 13))

print("x2")
ax3 = fig.add_subplot(lines, 3, 4)
ax3.scatter(x_2, x_2, c=colors)
# 为每个点添加序号标签
for i, txt in enumerate(x_1):
    ax3.text(x_2[i], x_2[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')

ax3 = fig.add_subplot(lines, 3, 5)
ax3.scatter(x_2_1, y_2_1, c=colors)
# 为每个点添加序号标签
for i, txt in enumerate(x_1):
    ax3.text(x_2_1[i], y_2_1[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')

ax3 = fig.add_subplot(lines, 3, 6)
ax3.scatter([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], edge_weight_for_graph1)
ax3.set_ylim((-0, 1))
ax3.set_xlim((0, 13))

print("x3")
ax3 = fig.add_subplot(lines, 3, 7)
ax3.scatter(x_3, y_3, c=colors)
# 为每个点添加序号标签
for i, txt in enumerate(x_1):
    ax3.text(x_3[i], y_3[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')

ax3 = fig.add_subplot(lines, 3, 8)
ax3.scatter(x_3_1, y_3_1, c=colors)
# 为每个点添加序号标签
for i, txt in enumerate(x_1):
    ax3.text(x_3_1[i], y_3_1[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')

ax3 = fig.add_subplot(lines, 3, 9)
ax3.scatter([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], edge_weight_for_graph2)
ax3.set_ylim((-0, 1))
ax3.set_xlim((0, 13))

# print("x4")
# ax3 = fig.add_subplot(lines, 2, 7)
# ax3.scatter(x_4, y_4, c=colors)
# for i, txt in enumerate(x_1):
#     ax3.text(x_4[i], y_4[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')
#
# ax3 = fig.add_subplot(lines, 2, 8)
# ax3.scatter([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], edge_weight_for_graph3)
# ax3.set_ylim((-0, 1))
# ax3.set_xlim((0, 13))

# print("x5")
# ax3 = fig.add_subplot(lines, 2, 9)
# ax3.scatter(x_5, y_5, c=colors)
# for i, txt in enumerate(x_1):
#     ax3.text(x_5[i], y_5[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')
#
# ax3 = fig.add_subplot(lines, 2, 10)
# ax3.scatter([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], edge_weight_for_graph4)
# ax3.set_ylim((-2, 2))
# ax3.set_xlim((0, 13))
#
# print("x6")
# ax3 = fig.add_subplot(lines, 2, 11)
# ax3.scatter(x_6, y_6, c=colors)
# for i, txt in enumerate(x_1):
#     ax3.text(x_6[i], y_6[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')
#
# ax3 = fig.add_subplot(lines, 2, 12)
# ax3.scatter([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], edge_weight_for_graph5)
# ax3.set_ylim((-2, 2))
# ax3.set_xlim((0, 13))

print("x7")
ax3 = fig.add_subplot(lines, 3, 10)
ax3.scatter(x_7, y_7, c=colors)
for i, txt in enumerate(x_1):
    ax3.text(x_7[i], y_7[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')

ax3 = fig.add_subplot(lines, 3, 11)
ax3.scatter(x_7_1, y_7_1, c=colors)
# 为每个点添加序号标签
for i, txt in enumerate(x_1):
    ax3.text(x_7_1[i], y_7_1[i], f'{i}', color='black', fontweight='bold', ha='center', va='center')

ax3 = fig.add_subplot(lines, 3, 12)
ax3.scatter([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0], edge_weight_for_graph6)
ax3.set_ylim((-0, 1))
ax3.set_xlim((0, 13))

fig.tight_layout(pad=2)
plt.show()
